import json
import torch
import torch.nn as nn


class NERBiLSTM(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, dropout, word2id, tag2id):
        super().__init__()

        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        # 加不加1都可以
        self.vocab_size = len(word2id) + 1
        self.tag_size = len(tag2id)
        self.tag2id = tag2id
        self.name = "BiLSTM"

        # embedd层
        self.embed = nn.Embedding(self.vocab_size, self.embedding_dim)
        # 随机失活层
        self.dropout = nn.Dropout(p=dropout)
        # BiLSTM
        self.lstm = nn.LSTM(self.embedding_dim, self.hidden_dim // 2,
                            bidirectional=True, batch_first=True)
        # 输出层
        self.out = nn.Linear(self.hidden_dim, self.tag_size)

    def forward(self, x, mask):
        """
        前向传播
        :param x:形状为[batch_size,seq_len]
        :param mask: 形状为[batch_size,seq_len]
        :return:
        """
        # embed_x=[batch_size,seq_len,embedding_dim]
        embed_x = self.embed(x)
        # output=[batch_size,seq_len,hidden_dim]
        output, hidden = self.lstm(embed_x)
        # 掩码
        output = output * mask.unsqueeze(-1)
        # dropout
        output = self.dropout(output)
        # result=[batch_size,seq_len,tag_size]
        result = self.out(output)
        return result


if __name__ == '__main__':

    ...
